from datam import *
from model import *
from train_lightning import *

if __name__ == '__main__':
    # data_mnist = DataM("../data/",32,0)
    data_mnist = DataM(config.data_dir, config.BATCH_SIZE, config.AVAIL_GPUS)
    data_mnist.setup()


    # model = Model.load_from_checkpoint("best.ckpt")
    # model.eval()
    # data_mnist = DataM("../data/",32,0)
    # data_mnist.setup()
    # data,label = next(iter(data_mnist.test_dataloader()))
    #
    # prediction = model(data)
    # print(prediction)



    trainer = Trainer(gpus=config.AVAIL_GPUS, max_epochs=config.max_epochs, logger=wandb_logger)
    model = Model.load_from_checkpoint("best.ckpt")
    result = trainer.test(model,data_mnist.test_dataloader())
    print(result)



